#![feature(test)]

extern crate test;

use criterion::{criterion_group, criterion_main, Criterion};
use protoactor_rs::actor::{Actor, ActorContext, ActorError, DefaultActorContext, PID, Props};
use protoactor_rs::remote::RemoteActorSystem;
use protoactor_rs::system::ActorSystem;
use std::sync::Arc;
use std::time::Duration;
use tokio::runtime::Runtime;

// Messages
#[derive(Clone)]
struct Ping;

#[derive(Clone)]
struct Pong;

#[derive(Clone)]
struct StartRemote {
    sender: PID,
}

// Echo Actor
struct EchoActor {
    sender: Option<PID>,
}

#[async_trait::async_trait]
impl Actor for EchoActor {
    type Context = DefaultActorContext;

    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        if let Some(start) = msg.downcast_ref::<StartRemote>() {
            println!("Starting echo actor for {}", start.sender);
            self.sender = Some(start.sender.clone());
            ctx.respond(Box::new(()));
        } else if msg.downcast_ref::<Ping>().is_some() {
            if let Some(sender) = &self.sender {
                ctx.send(sender, Box::new(Pong)).await?;
            }
        }
        Ok(())
    }
}

// Benchmark Node 1 (Sender)
async fn setup_node1(port: u16) -> (RemoteActorSystem, PID) {
    let mut system = RemoteActorSystem::new(&format!("127.0.0.1:{}", port)).unwrap();
    system.start().await.unwrap();
    
    let props = Props::from_producer(|| EchoActor { sender: None });
    let pid = system.spawn_local("sender", props).await.unwrap();
    
    (system, pid)
}

// Benchmark Node 2 (Receiver)
async fn setup_node2(port: u16) -> RemoteActorSystem {
    let mut system = RemoteActorSystem::new(&format!("127.0.0.1:{}", port)).unwrap();
    
    let props = Props::from_producer(|| EchoActor { sender: None });
    system.register("echo", props);
    
    system.start().await.unwrap();
    system
}

fn remote_pingpong_benchmark(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("remote_pingpong", |b| {
        b.iter_custom(|iters| {
            rt.block_on(async {
                // Setup nodes
                let (node1, sender_pid) = setup_node1(12000).await;
                let node2 = setup_node2(12001).await;
                
                // Get remote actor reference
                let remote_pid = node1.spawn_remote(
                    "127.0.0.1:12001",
                    "echo",
                    "remote",
                    Duration::from_secs(5)
                ).await.unwrap();
                
                // Start remote actor
                node1.request::<()>(
                    &remote_pid,
                    StartRemote { sender: sender_pid.clone() },
                    Duration::from_secs(5)
                ).await.unwrap();
                
                // Benchmark loop
                let start = std::time::Instant::now();
                for _ in 0..iters {
                    node1.send(&remote_pid, Ping).await.unwrap();
                    node1.request::<Pong>(&sender_pid, Ping, Duration::from_secs(1)).await.unwrap();
                }
                let duration = start.elapsed();
                
                // Cleanup
                node1.stop().await.unwrap();
                node2.stop().await.unwrap();
                
                duration
            })
        });
    });
}

fn local_pingpong_benchmark(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("local_pingpong", |b| {
        b.iter_custom(|iters| {
            rt.block_on(async {
                let system = ActorSystem::new().unwrap();
                
                // Create actors
                let props = Props::from_producer(|| EchoActor { sender: None });
                let sender_pid = system.spawn::<EchoActor>(props.clone()).await.unwrap();
                let echo_pid = system.spawn::<EchoActor>(props).await.unwrap();
                
                // Start echo actor
                system.request::<()>(
                    &echo_pid,
                    StartRemote { sender: sender_pid.clone() },
                    Duration::from_secs(5)
                ).await.unwrap();
                
                // Benchmark loop
                let start = std::time::Instant::now();
                for _ in 0..iters {
                    system.send(&echo_pid, Ping).await.unwrap();
                    system.request::<Pong>(&sender_pid, Ping, Duration::from_secs(1)).await.unwrap();
                }
                
                start.elapsed()
            })
        });
    });
}

fn skynet_benchmark(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    
    c.bench_function("skynet", |b| {
        b.iter_custom(|_iters| {
            rt.block_on(async {
                const NUM_ACTORS: usize = 1_000_000;
                let system = ActorSystem::new().unwrap();
                
                // Create a large number of actors
                let mut pids = Vec::with_capacity(NUM_ACTORS);
                let props = Props::from_producer(|| EchoActor { sender: None });
                
                let start = std::time::Instant::now();
                for i in 0..NUM_ACTORS {
                    let pid = system.spawn::<EchoActor>(props.clone()).await.unwrap();
                    pids.push(pid);
                }
                
                // Send a message to each actor
                for pid in &pids {
                    system.send(pid, Ping).await.unwrap();
                }
                
                start.elapsed()
            })
        });
    });
}

criterion_group!(benches, remote_pingpong_benchmark, local_pingpong_benchmark, skynet_benchmark);
criterion_main!(benches); 